Make SD3Transformer2DModel hidden states contiguous#14186
Conversation
| hidden_states = hidden_states.contiguous() | ||
| if encoder_hidden_states is not None: | ||
| encoder_hidden_states = encoder_hidden_states.contiguous() |
There was a problem hiding this comment.
Could also comment on the position in the block of code where contiguous() should be applied?
There was a problem hiding this comment.
You're right. After investigating, the non-contiguous tensor originates from PatchEmbed.forward - flatten(2).transpose(1, 2) (BCHW → BNC) produces a non-contiguous layout.
diffusers/src/diffusers/models/embeddings.py
Line 562 in 4f37263
So the fix is more precisely placed right after self.pos_embed(), not inside the loop.
Updated the patch: moved .contiguous() to immediately after self.pos_embed(hidden_states), dropped the encoder_hidden_states call, and added a comment explaining the root cause.
There was a problem hiding this comment.
Did you verify the performance table reported in the original description of this PR?
There was a problem hiding this comment.
Yes, the numbers in the description have been updated. The re-run results are consistent with the previously reported measurements, with no significant deviation.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
Very cool work. Would be nice to also do this for others? |
I'm not sure if other models have the same issue, but will follow up and create new PR if I find cases where the similar fix applies. |
3 participants
What does this PR do?
Ensures hidden states are contiguous after each
JointTransformerBlockcall inSD3Transformer2DModel.On ROCm, non-contiguous tensors produced by
JointTransformerBlockcause performance degradation that accumulates across transformer blocks.Benchmarks (SD3-medium, 512×512, 28 steps, fp16, No text_encoder_3):
Before submitting
.ai/review-rules.md?documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
@sayakpaul @yiyixuxu
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.